package main

import (
	"fmt"
	"github.com/throttled/throttled"
	"github.com/throttled/throttled/store/memstore"
	"log"
	"net/http"
)

// throttled 是一款轻量的限流组件，可以非常轻松地实现单机或集群的 QPS 限流能力。
// go get -v github.com/throttled/throttled
// go get -v github.com/hashicorp/golang-lru
func main() {
	// 在 store 中添加 key 的数量限制
	//并内置了 memstore，redis store 等元信息存储实现，我们可以根据具体的使用场景实现单机限流和集群限流。
	store, err := memstore.New(65536)
	if err != nil {
		log.Fatal(err)
	}

	// 配置限流规则
	// 按请求 path 进行限流的 http 服务，允许同一个 path 每分钟 20 次请求，并支持短时间内最多 5 个请求的 burs
	quota := throttled.RateQuota{
		MaxRate:  throttled.PerMin(20),
		MaxBurst: 5,
	}
	rateLimiter, err := throttled.NewGCRARateLimiter(store, quota)
	if err != nil {
		log.Fatal(err)
	}

	httpRateLimiter := throttled.HTTPRateLimiter{
		RateLimiter: rateLimiter,
		// 根据 path 进行限流
		VaryBy: &throttled.VaryBy{Path: true},
		// 自定义处理规则
		DeniedHandler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("被限流了\n"))
		}),
	}

	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, "hello, world\n")
	})
	http.ListenAndServe(":8080", httpRateLimiter.RateLimit(handler))
}
